#######################################################################################
# File: BSMLM.R                                                             
# Author: Xun Pang, Stephen Chaudoin, Helen Milner                                                             
# Creation Date: 10/11/2011                                                           

 #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@
 ## MCMC implementation of a TSCS model with spatial lags within the MLM structure 
 ##
 ## \imput Data are fitted as \y_i's
 ## Y--- outcome data, vector
 ## X--- covariates with fixed effects
 ## Si--- covariates with unit-specific effects
 ## Dt--- covariates with time-specific effects
 ## Ai--- predictor to explain the variations of beta_s
 ## Zt--- predictor to explain the variations of beta_d
 ## W---weight matrix,  a list of matrices, can be just a single matrix, but in a list form
 ## WT--- vector of indicator of which W used in a specific time period. A vector of T length.if
 ##       a time invariant W is used, then the vector is all 1's; the unique values of the vector
 ##       should equal to the length of the list of W
 ##
 ## \para
 ##           
 ## group 1 -- priors
 ##          beta0: prior mean of beta, a K \time 1 matrix
 ##          B0: prior covariance matrix of beta, a K \times K matrix
 ##          a0 and b0: degrees of freedom (d0 > rup) and scale matrix (rup \times rup) of the precision
 ##
 ##
 ## group 6 -- storage 
 ##        Dstore, Estore, betastore, bistore, rhostore, ctstore, logmarglike
 ##        (note: no matter whether a parameter is monitored, all the matrices have to be specified 
 ##         and passed into the function. User does not have to specify those matrices)      
 #@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@@*

 library(msm)

 BSMLMWT.MCMC <- function(Uid,Tid, Y, X1, Si, Dt, Ai, Zt,W, WT, timeint.add=TRUE, unitint.add=TRUE,
                       rho, precision, beta0, B0, a0, b0,  D0, d0, E0, e0,
                       mcmc, burnin, thin, tracking){
   ################################
   # Form Data in Reduced Form
   ################################
   Uid <- Index(Uid)
   Tid <- Index(Tid)
   DataReduce <- RFGenerator(Y,X1, Si, Dt, Ai, Zt, Uid, Tid,
                            timeint.add, unitint.add)
   Uid <- Index(DataReduce[[1]])
   Tid <- Index(DataReduce[[2]])
   Y <- DataReduce[[3]]
   X <- DataReduce[[4]]
   S <- DataReduce[[5]]   # interaction terms in the reduced form
   F <- DataReduce[[6]]   # interaction terms in the reduced form
   # get the dimension
   K <- ncol(X)
   cat("K=", K, "\n")
   n <- length(Y)
   rup <- ncol(S)
   rtp <- ncol(F)
   entry <- 1:n
   N <- length(unique(Uid))
   T <- length(unique(Tid))
   RUP <- ncol(Ai)    # number of the second level predictors
   RTP <- ncol(Zt)    # number of the second level predictors   

     # storage
     after.iter <- mcmc * thin
     total.iter <- burnin + after.iter
     total.store <- mcmc
     rho.store <- matrix(NA, ncol=T, nrow=mcmc)
     beta.store <- matrix(NA, ncol=K, nrow=mcmc)
     precision.store <- rep(NA, mcmc)
     bi.store <- matrix(NA, nrow=mcmc, ncol=rup*N)
     ct.store <- matrix(NA, nrow=mcmc, ncol=rtp*T)
     
     # Do not store D and E
     b.mat <- matrix(0, nrow=N, ncol=rup)
     c.mat <- matrix(0, nrow=T, ncol=rtp)
   
     # inner product of X
     x.mat <- t(X)%*%X
     # rearrange X in order of \y_t
     Xt.mat <- matrix(0, ncol=K, nrow=1)
     Yt <- c()
     TTid <- c()   # create index when the data are arranged according to time
     UUid <- c()   # new idex of the i's
     trans.Y <- c()
     for (t in 1: T){
      rows <- entry[which(Tid==t)]
      Xtt <- X[rows,]
      Xt.mat <- rbind(Xt.mat, Xtt)
      ytm <- Y[rows]
      Yt <- c(Yt, ytm)
      Wc <- as.matrix(W[[WT[t]]])
      trans.yt <- Wc%*%ytm
      trans.Y <- c(trans.Y, trans.yt)
      TTid <- c(TTid, rep(t, length(rows)))
      UUid <- c(UUid, Uid[rows])
     }
     Xt.mat <- Xt.mat[-1,]
     B0.inv <- solve(B0)
     priorb <- B0.inv%*%beta0 
   #######################
   # MCMC Updating Begins
   #######################  

   #Tracking the process of MCMC simulation

   for(g in 1:(total.iter)) {
      if (g%% tracking ==0) cat("g=",g,"\n")

      # MCMC updating 1: the variance-covariance matrices

       Eupdate <- Covariance.update(Covariate=F, prior.freedom=e0, prior.scale=E0,
                              dimension=rtp, obs.length=T, current.value=c.mat)
       E <- Eupdate[[1]]
       Dupdate <- Covariance.update(Covariate=S, prior.freedom=d0, prior.scale=D0,
                              dimension=rup, obs.length=N, current.value=b.mat)
       D <- Dupdate[[1]]

       ## MCMC Updating 2: beta
       # loop begins
       raniw <- c()
       for (i in 1:N){
         rowsi <- entry[which(Uid==i)]
         Si <- as.matrix(S[rowsi,])
         rani <- Si%*% b.mat[i,]
         raniw <- c(raniw, rani)
       }
       rwy <- c()  # rwy is the transformed y arrange according to time
       rantw <- c() # Fct as \y_t
       raniwt <- c()  # Sbi as \y_t
       for (t in 1: T){                              
         rows <- entry[which(Tid==t)]               
         Ft <- as.matrix(F[rows,])
         I <- diag(length(rows))
         Wc <- as.matrix(W[[WT[t]]])
         A <- I-rho[t]*Wc
         ytm <- Y[rows]
         rwyt <-  A%*%ytm
         rwy <- c(rwy, rwyt)
         rant <- Ft%*%c.mat[t,]
         rantw <- c(rantw, rant)
         raniwt <- c(raniwt, raniw[rows])
       }
       B1 <- solve(B0.inv + precision* x.mat)
       res.beta <- rwy-raniwt-rantw
       beta1 <- B1%*%(priorb+ precision*t(Xt.mat)%*% res.beta)
       beta <- mvrnorm(1, beta1, B1)
    
       # MCMC Updating 2: precision
       x.mean <- Xt.mat%*%beta
       a1 <- (a0+n)/2
       b1 <- (b0+ sum((res.beta-x.mean)^2))/2
       precision <- rgamma(1, a1, scale=1/b1)
       #cat("precision=",precision,"\n")
      
       # MCMC Updating 3: rho
       who.res <- x.mean + rantw + raniwt
       transy <- c()  #transformed y rW\y_t
       for(t in 1:T){
         rows <- entry[which(TTid==t)]
         yt <- Yt[rows]
         wytt <- trans.Y[rows]
         mean.x <- who.res[rows]
         Phi1 <- solve(precision* sum(wytt^2))
         phi1 <- Phi1* precision* sum(t(wytt)%*%(yt-mean.x))
         # Normally, we do not need the lower bound
         rho.p <- rtnorm(1, phi1, sqrt(Phi1), upper=1) # truncated draw
         # rho.p <- rtnorm(1, phi1, sqrt(Phi1), upper=1, lower=-1) # truncated draw
         #rho.p <- rtnorm(1, rho[t], 0.01, upper=1, lower=-1) #random walk chain
         # A-R
         Wc <- as.matrix(W[[WT[t]]])       
         I <- diag(length(rows))
         A.new <- I-rho.p*Wc
         A.old <- I-rho[t]*Wc
         A.p <- det(A.new)
         A.c <- det(A.old)
         deter.new <-  log(abs(A.p))
         deter.old <-  log(abs(A.c))
         res.new <- sum(A.new%*%yt-mean.x)^2
         res.old <- sum(A.old%*%yt-mean.x)^2
         alpha <- deter.new-deter.old-precision/2*(res.new-res.old)
         if (alpha > 0) rho[t] <- rho.p
         transy <- c(transy, rho[t]*wytt)
       }
       #   cat("alpha=",alpha,"\n")

       # MCMC updating: {ct}
       res.part <- Yt-transy-x.mean
       res.c <- res.part-raniwt
       newrantw <- c()
       for (t in 1:T){
         rowst <- entry[which(Tid==t)]
         Ft <- as.matrix(F[rowst,])
         rowst2 <- entry[which(TTid==t)]
         res.ct <- res.c[rowst2,]
         Ct1 <- solve(E+precision*t(Ft)%*% Ft)
         ct1 <- Ct1%*%(precision* t(Ft)%*%res.ct)
         newc <- mvrnorm(1, ct1, Ct1)
         c.mat[t,] <- newc
         newrant <- Ft%*%newc
         newrantw <- c(newrantw, newrant)
       }

       # MCMC updating: {bi}
       res.bt <- res.part-newrantw
       for (i in 1:N){
         rows <- entry[which(UUid==i)]
         rows2 <- entry[which(Uid==i)]
         Si <- as.matrix(S[rows2,])
         res.bi <- res.bt[rows]
         Bi1 <- solve(D+precision*t(Si)%*%Si)
         bi1 <- Bi1%*%(precision*t(Si)%*%res.bi)
         newb <- mvrnorm(1, bi1, Bi1)
         b.mat[i,] <- newb
       }
      # store
       gab <- g - burnin
       if (gab > 0){
       if ((gab%%thin ==0)| gab == after.iter){
          gg <- gab/thin  
          beta.store[gg,] <- beta
          rho.store[gg,] <- rho
          precision.store[gg] <- precision
          bi.store[gg,] <- as.vector(b.mat)
          ct.store[gg,] <- as.vector(c.mat)
        }
      }
     }

    # return
     mcmc.output <- list(beta=beta.store, rho=rho.store, precision=precision.store,
                         bi=bi.store, ct=ct.store)
     return(mcmc.output)
 }
        
        

